import argparse
import yaml
import os


def parse_args():
    parser = argparse.ArgumentParser(description="ViTAD Training")
    parser.add_argument('--config', type=str, required=True, help="Path to the YAML configuration file.")
    parser.add_argument('--lr', type=float, default=None, help="Learning rate override.")
    parser.add_argument('--weight_decay', type=float, default=None, help="Weight decay override.")
    args = parser.parse_args()

    # Load YAML configuration
    with open(args.config, 'r') as f:
        cfg = yaml.safe_load(f)

    # Override with command-line arguments if provided
    if args.lr is not None:
        cfg['trainer']['lr'] = args.lr
    if args.weight_decay is not None:
        cfg['trainer']['weight_decay'] = args.weight_decay

    return args, cfg
